from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.chains import LLMChain
from prometheus_api_client import PrometheusConnect
from prometheus_api_client import MetricSnapshotDataFrame, MetricRangeDataFrame
import numpy as np
import matplotlib.pyplot as plt
import base64
from io import BytesIO, StringIO
from datetime import datetime, timedelta
from config import config

class PrometheusQueryService:
    def __init__(self):
        """初始化Prometheus查询服务"""
        # 初始化LangChain
        self.llm = ChatOpenAI(
            model_name=config.LANGCHAIN_MODEL,
            temperature=config.LANGCHAIN_TEMPERATURE
        )
        
        # 初始化Prometheus连接
        self.prometheus = PrometheusConnect(
            url=config.PROMETHEUS_URL,
            disable_ssl=True  # 生产环境中应启用SSL
        )
        
        # 初始化提示模板
        self._init_prompt_template()
    
    def _init_prompt_template(self):
        """初始化提示模板"""
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", """你是一个PromQL专家，能够将自然语言查询转换为正确的PromQL查询语句。
            请根据以下规则生成PromQL：
            
            1. 只返回PromQL查询语句，不要包含任何解释或额外信息。
            2. 使用正确的PromQL语法和函数。
            3. 如果需要时间范围，请使用相对时间，如5m、1h、1d等。
            4. 如果用户没有指定时间范围，默认使用1小时。
            5. 确保查询是有效的，可以直接在Prometheus中执行。
            
            例如：
            自然语言："过去5分钟内api_server的请求量"
            PromQL：sum(rate(http_requests_total{job="api_server"}[5m]))
            
            自然语言："过去1小时内每个服务的错误率"
            PromQL：sum(rate(http_requests_total{status=~"5.."}[1h])) by (service) / sum(rate(http_requests_total[1h])) by (service)
            """),
            ("human", "{query}")
        ])
        
        self.chain = LLMChain(llm=self.llm, prompt=self.prompt)
    
    def natural_language_to_promql(self, natural_language_query: str) -> str:
        """
        将自然语言查询转换为PromQL
        
        Args:
            natural_language_query: 用户输入的自然语言查询
            
        Returns:
            str: 生成的PromQL查询语句
        """
        try:
            result = self.chain.run(natural_language_query)
            return result.strip()
        except Exception as e:
            raise Exception(f"Failed to convert natural language to PromQL: {str(e)}")
    
    def query_prometheus(self, promql: str, time_range: str = "1h") -> dict:
        """
        执行PromQL查询并返回结果
        
        Args:
            promql: PromQL查询语句
            time_range: 时间范围，默认为1小时
            
        Returns:
            dict: 查询结果
        """
        try:
            # 解析时间范围
            end_time = datetime.now()
            start_time = end_time - self._parse_time_range(time_range)
            
            # 执行查询
            result = self.prometheus.custom_query_range(
                query=promql,
                start_time=start_time,
                end_time=end_time,
                step="1m"  # 1分钟步长
            )
            
            return result
        except Exception as e:
            raise Exception(f"Failed to query Prometheus: {str(e)}")
    
    def _parse_time_range(self, time_range: str) -> timedelta:
        """
        解析时间范围字符串
        
        Args:
            time_range: 时间范围字符串，如"5m"、"1h"、"1d"
            
        Returns:
            timedelta: 解析后的的时间范围
        """
        units_map = {
            's': 'seconds',
            'm': 'minutes',
            'h': 'hours',
            'd': 'days',
            'w': 'weeks'
        }
        
        try:
            value = int(time_range[:-1])
            unit = time_range[-1]
            
            if unit not in unit_map:
                raise ValueError(f"Unsupported time unit: {unit}")
            
            return timedelta(**{unit_map[unit]: value})
        except (ValueError, IndexError):
            # 默认返回1小时
            return timedelta(hours=1)
    
    def generate_plot(self, data: dict, graph_type: str = config.DEFAULT_GRAPH_TYPE) -> str:
        """
        生成图表并返回base64编码
        
        Args:
            data: Prometheus查询结果
            graph_type: 图表类型，支持"line"、"bar"、"area"
            
        Returns:
            str: base64编码的图表
        """
        if not data:
            return ""
        
        try:
            # 设置中文字体
            plt.rcParams['font.sans-serif'] = ['WenQuanYi Zen Hei']
            plt.rcParams['axes.unicode_minus'] = False
            
            # 创建图表
            fig, ax = plt.subplots(figsize=(config.GRAPH_WIDTH/100, config.GRAPH_HEIGHT/100))
            
            # 处理数据
            for metric in data:
                metric_name = metric['metric'].get('__name__', 'unknown_metric')
                labels = metric['metric']
                
                # 创建标签字符串
                label_parts = []
                for key, value in labels.items():
                    if key != '__name__':
                        label_parts.append(f"{key}={value}")
                label = f"{metric_name} ({', '.join(label_parts)})" if label_parts else metric_name
                
                # 提取数据点
                timestamps = [point[0] for point in metric['values']]
                values = [float(point[1]) for point in metric['values']]
                
                # 转换时间戳
                times = [datetime.fromtimestamp(ts) for ts in timestamps]
                
                # 绘制图表
                if graph_type == 'line':
                    ax.plot.plot(times, values, label=label)
                elif graph_type == 'bar':
                    ax.bar(times, values, label=label, alpha=0.7)
                elif graph_type == 'area':
                    ax.fill_between(times, values, alpha=0.3, label=label)
                    ax.plot(times, values, label=label)
            
            # 设置图表属性
            ax.set_title('Prometheus 监控数据', fontsize=16)
            ax.set_xlabel('时间', fontsize=12)
            ax.set_ylabel('值', fontsize=12)
            ax.legend(loc='best')
            ax.grid(True, alpha=0.3)
            
            # 自动调整布局
            plt.tight_layout()
            
            # 保存图表到内存
            buffer = BytesIO()
            plt.savefig(buffer, format='png', dpi=300, bbox_inches='tight')
            buffer.seek(0)
            
            # 转换为base64
            image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
            
            # 关闭图表
            plt.close()
            
            return image_base64
        except Exception as e:
            print(f"Failed to generate plot: {str(e)}")
            return ""
    
    def query_with_natural_language(self, natural_language_query: str, time_range: str = "1h", graph_type: str = config.DEFAULT_GRAPH_TYPE) -> dict:
        """
        使用自然语言查询Prometheus并返回结果
        
        Args:
            natural_language_query: 自然语言查询
            time_range: 时间范围
            graph_type: 图表类型
            
        Returns:
            dict: 查询结果，包含promql、data和plot
        """
        try:
            # 转换为PromQL
            promql = self.natural_language_to_promql(natural_language_query)
            
            # 查询Prometheus
            data = self.query_prometheus(promql, time_range)
            
            # 生成图表
            plot = self.generate_plot(data, graph_type)
            
            return {
                'success': True,
                'promql': promql,
                'data': data,
                'plot': plot
            }
        except Exception as e:
            return {
                'success': False,
                'error': str(e)
            }
